package skewreduce.framework.planner;

import java.io.BufferedOutputStream;
import java.io.BufferedReader;
import java.io.DataInputStream;
import java.io.DataOutputStream;
import java.io.EOFException;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.FileReader;
import java.io.IOException;
import java.io.PrintStream;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.IdentityHashMap;
import java.util.List;
import java.util.PriorityQueue;
import java.util.Set;
import java.util.SortedSet;
import java.util.TreeSet;

import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.FSDataInputStream;
import org.apache.hadoop.fs.FileStatus;
import org.apache.hadoop.fs.FileSystem;
import org.apache.hadoop.fs.Path;
import org.apache.hadoop.mapred.ClusterStatus;
import org.apache.hadoop.mapred.JobClient;
import org.apache.hadoop.mapred.JobConf;
import org.apache.hadoop.util.GenericOptionsParser;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import skewreduce.framework.Schedule;
import skewreduce.framework.logical.LogicalOp;
import skewreduce.framework.physical.PhysicalOp;
import skewreduce.framework.planner.partition.Partitioner;
import skewreduce.lib.ITuple;
import skewreduce.lib.KahanSum;
import skewreduce.lib.Partition;


/**
 * using sample and known execution, fit constants for different cost model.
 * @author yongchul
 *
 */
public class FitCostTime {
	private static final Logger LOG = LoggerFactory.getLogger(FitCostTime.class);
	
	final Configuration conf;
	final FileSystem fs;
	List<ITuple> sample = new ArrayList<ITuple>();
	
	private IdentityHashMap<PartitionNode,PhysicalOp.ID> node2id = new IdentityHashMap<PartitionNode,PhysicalOp.ID>();
	private HashMap<PhysicalOp.ID,Double> id2time = new HashMap<PhysicalOp.ID,Double>();
	
	class LinearRegression {
		ArrayList<Double> x; 
		ArrayList<Double> y;

		KahanSum sumx2 = new KahanSum();
		KahanSum sumx = new KahanSum();
		KahanSum sumy = new KahanSum();
		
		KahanSum xxbar = new KahanSum();
		KahanSum yybar = new KahanSum();
		KahanSum xybar = new KahanSum();
		
		double xbar;
		double ybar;
		
		double beta0;
		double beta1;
		
        KahanSum rss = new KahanSum();      // residual sum of squares
        KahanSum ssr = new KahanSum();      // regression sum of squares
		
		LinearRegression() {
			this(10);
		}
		
		LinearRegression(int cap) {
			x = new ArrayList<Double>(cap);
			y = new ArrayList<Double>(cap);
		}
		
		public void reset() {
			x.clear(); y.clear();
			sumx.reset();
			sumy.reset();
			sumx2.reset();
			xxbar.reset();
			yybar.reset();
			xybar.reset();
		}
		public void update(double x,double y) {
			this.x.add(x);
			this.y.add(y);
			
			sumx.add(x);
			sumy.add(y);
			sumx2.add(x*x);
		}
		
		public void doFinal() {
			xbar = sumx.value() / x.size();
			ybar = sumy.value() / x.size();
			
			for ( int i = 0; i < x.size(); ++i ) {
				double xx = x.get(i);
				double yy = y.get(i);
				
				xxbar.add((xx - xbar)*(xx - xbar));
				yybar.add((yy - ybar)*(yy - ybar));
				xybar.add((xx - xbar)*(yy - ybar));
			}
			
			beta1 = xybar.value() / xxbar.value();
			beta0 = ybar - beta1 * xbar;
			
	        for (int i = 0; i < x.size(); i++) {
	            double fit = beta1*x.get(i) + beta0;
	            rss.add((fit - y.get(i)) * (fit - y.get(i)));
	            ssr.add((fit - ybar) * (fit - ybar));
	        }
		}
		
		public double getBeta0() { return beta0; }
		public double getBeta1() { return beta1; }
		public double getR2() { return ssr.value() / yybar.value(); }
				
		@Override
		public String toString() {
			return String.format("y = %g * x + %g (R^2 = %g)",beta1,beta0,getR2());
		}

        // analyze results
//        int df = n - 2;
//        double svar  = rss / df;
//        double svar1 = svar / xxbar;
//        double svar0 = svar/n + xbar*xbar*svar1;
//        System.out.println("R^2                 = " + R2);
//        System.out.println("std error of beta_1 = " + Math.sqrt(svar1));
//        System.out.println("std error of beta_0 = " + Math.sqrt(svar0));
//        svar0 = svar * sumx2 / (n * xxbar);
//        System.out.println("std error of beta_0 = " + Math.sqrt(svar0));
//
//        System.out.println("SSTO = " + yybar);
//        System.out.println("SSE  = " + rss);
//        System.out.println("SSR  = " + ssr);
	}

	private EaggCostModel model;
	private PartitionNode rootNode;
	private Partition rootPartition;
	private double sampleRate;
	private double bandWidth;
	
	private static final int DEFAULT_SAMPLE_SIZE = 10;
	private int sampleWorkSize = DEFAULT_SAMPLE_SIZE;
	private int sampleMergeSize = DEFAULT_SAMPLE_SIZE;

	LinearRegression regWork = new LinearRegression();
	LinearRegression regMerge = new LinearRegression();
	
	private boolean verbose;
	
	private void partition(PartitionNode node) {
		ArrayList<ITuple> ls = new ArrayList<ITuple>();
		ArrayList<ITuple> rs = new ArrayList<ITuple>();
		
		Partition parent = node.getPartition();
		
		for ( ITuple t : node.getSample() ) {
			int child = parent.whichChild(t);
			if ( child == 0 ) ls.add(t);
			else rs.add(t);
		}
		
		ls.trimToSize();
		rs.trimToSize();
		
		node.setChildren(new PartitionNode(parent.getLeft(),ls),new PartitionNode(parent.getRight(),rs));
	}
	
	private void buildPartitionTree(PartitionNode node) {
		if ( node == null ) return;
		
		if ( ! node.getPartition().isLeaf() ) {
			partition(node);
			
			buildPartitionTree(node.getLeft());
			buildPartitionTree(node.getRight());
		}

		node2id.put(node,new PhysicalOp.ID(node.getPartition()));
	}
	
	private void fit(PartitionNode node) {
		if ( node.isLeaf() ) {
			// estimate cost
			double cost = model.costWork(node);
			double time = id2time.get( node2id.get(node) );
			
			// fit const
			if ( sampleWorkSize-- > 0 ) {
				regWork.update(cost,time);
			}
		} else {
			fit(node.getLeft());
			fit(node.getRight());
			
			double cost = model.costMerge(node,node.getLeft(),node.getRight());
			double time = id2time.get(node2id.get(node));
			
			// fit cost
			if ( sampleMergeSize-- > 0 )
				regMerge.update(cost, time);
		}
	}
	
	private void fit() {
		fit(rootNode);
		
		regWork.doFinal();
		regMerge.doFinal();
	}
	
	private void buildPartitionTree() {
		rootNode = new PartitionNode(rootPartition,sample);
		buildPartitionTree(rootNode);
	}
	
	//////////////////////////////
	
	FitCostTime(Configuration cnf) throws IOException {
		conf = cnf;
		fs = FileSystem.get(this.conf);
	}
	
	private int loadSample(Path fn) throws IOException {
		int before = sample.size();
		FSDataInputStream in = null;
		ITuple t = null;
		try {
			in = fs.open(fn);
			while ( (t = model.createTuple(in)) != null ) {
				sample.add(t);
			}
		} catch ( EOFException ignore ) {
			// we simply ignore this exception and throw all other IOExceptions
		} finally {
			if ( in != null ) try { in.close(); } catch ( IOException ignore ) {}
			in = null;
		}
		return sample.size() - before;
	}
	
	private void doMain(String[] args) throws Exception {
		boolean bad = false;
		String planFile = null;
		String timeFile = null;
		int i = 0;
		
		for ( ; i < args.length; ++i ) {
			if ( args[i].charAt(0) != '-' )
				break;
			
			if ( "-sampleRate".equals(args[i]) ) {
				sampleRate = Double.parseDouble(args[++i]);
				conf.setFloat("skewreduce.planner.sampleRate", (float)sampleRate);
			} else if ( "-model".equals(args[i]) ) {
				conf.set(EaggCostModel.MODEL_CLASS_PROP, args[++i]);
			} else if ( "-bandWidth".equals(args[i]) ) {
				bandWidth = Double.parseDouble(args[++i]);
				conf.setFloat("skewreduce.planner.partition.bandWidth", (float)bandWidth);
			} else if ( "-verbose".equals(args[i]) || "-v".equals(args[i]) ) {
				verbose = true;
			} else if ( "-plan".equals(args[i]) ) {
				planFile = args[++i];
			} else if ( "-time".equals(args[i]) ) {
				timeFile = args[++i];
			} else if ( "-sampleSize".equals(args[i]) ) {
				int sz = Integer.parseInt(args[++i]);
				this.sampleWorkSize = this.sampleMergeSize = sz;
			} else {
				System.err.println("Unknown option: "+args[i]);
				bad = true;
			}
		}
		
		model = EaggCostModel.getInstance(conf);
		
		bad = bad || model == null;

		if ( ! bad ) {
			for ( ; i < args.length; ++i ) {
				FileStatus[] fns = fs.globStatus(new Path(args[i]));
				for ( FileStatus fn : fns ) {
					int nsample = loadSample(fn.getPath());
					LOG.info("{} samples have loaded from {}",nsample,fn.getPath());
	
				}
			}
		}
		
		LOG.info("Total {} samples have been loaded",sample.size());
		
		if ( bad || sample.isEmpty() ) {
			// FIXME print usage information
			System.exit(-1);
		}
		
		rootPartition = model.getRootPartition();
		{
			// load plan
			DataInputStream input = new DataInputStream(new FileInputStream(planFile));
			rootPartition.readFields(input);
			input.close();
		}
		
		{
			// load timing
			BufferedReader input = new BufferedReader(new FileReader(timeFile));
			String line;
			while ( (line = input.readLine()) != null ) {
				String[] fields = line.split(",");
				PhysicalOp.ID id = new PhysicalOp.ID(fields[0]);
				if ( id.getType() != LogicalOp.Type.WORK && id.getType() != LogicalOp.Type.MERGE ) continue;
				double t = Double.parseDouble(fields[1]);
				id2time.put(id, t);
			}
			
			input.close();
		}
		
		buildPartitionTree();
		
		fit(); // bottom-up evaluation
		
		System.out.println("Work Estimate: "+ this.regWork);
		System.out.println("Merge Estimate: "+ this.regMerge);
	}

	/**
	 * @param args
	 */
	public static void main(String[] args) throws Exception {
        GenericOptionsParser options = new GenericOptionsParser(new Configuration(),args);
        Configuration conf = options.getConfiguration();
        String[] rest = options.getRemainingArgs();
        
        //LOG.info("current working directory = {}",System.getProperty("user.dir"));

        if ( rest == null || rest.length == 0 ) {
            System.err.println("No class argument is given!");
            System.exit(1);
        }
        
        FitCostTime prog = new FitCostTime(conf);
        prog.doMain(rest);		
	}

}
